#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
情感指数计算工具
用于将情感标签转换为数值并计算综合情感指数
"""

import pandas as pd
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
import os


def map_sentiment_to_value(sentiment):
    """
    将情感标签映射为数值
    
    参数:
        sentiment: 情感标签，可以是"正面"、"中性"、"负面"
    
    返回:
        数值化的情感得分: 正面=1, 中性=0, 负面=-1
    """
    sentiment_map = {
        "正面": 1.0,
        "中性": 0.0,
        "负面": -1.0
    }
    
    # 如果情感标签不在映射表中，则返回0（中性）
    return sentiment_map.get(sentiment, 0.0)


def calculate_weight(row, click_weight=1.0, comment_weight=2.0, majia_penalty=0.5, 
                    min_click=0, max_click=1000, min_comment=0, max_comment=100):
    """
    计算每条评论的权重
    
    参数:
        row: DataFrame的一行数据
        click_weight: 点击量权重系数
        comment_weight: 评论数权重系数
        majia_penalty: 马甲账号的权重惩罚系数
        min_click/max_click: 点击量的最小/最大归一化范围
        min_comment/max_comment: 评论数的最小/最大归一化范围
    
    返回:
        该条评论的权重
    """
    # 基础权重为1
    weight = 1.0
    
    # 如果存在点击数，增加权重
    if 'post_click_count' in row and pd.notna(row['post_click_count']):
        click_count = max(min_click, min(max_click, row['post_click_count']))
        # 归一化点击数并计算点击权重
        normalized_click = (click_count - min_click) / (max_click - min_click) if max_click > min_click else 0
        weight += normalized_click * click_weight
    
    # 如果存在评论数，增加权重
    if 'post_comment_count' in row and pd.notna(row['post_comment_count']):
        comment_count = max(min_comment, min(max_comment, row['post_comment_count']))
        # 归一化评论数并计算评论权重
        normalized_comment = (comment_count - min_comment) / (max_comment - min_comment) if max_comment > min_comment else 0
        weight += normalized_comment * comment_weight
    
    # 如果是马甲账号，降低权重
    if 'user_is_majia' in row and pd.notna(row['user_is_majia']):
        if row['user_is_majia'] == True or row['user_is_majia'] == 1:
            weight *= majia_penalty
    
    return weight


def calculate_daily_sentiment_index(df, date_column='post_time', sentiment_column='sentiment',
                                   click_column='post_click_count', comment_column='post_comment_count',
                                   majia_column='user_is_majia', click_weight=1.0, comment_weight=2.0,
                                   majia_penalty=0.5, normalize=True):
    """
    按天计算情感指数
    
    参数:
        df: 包含评论数据的DataFrame
        date_column: 日期列名
        sentiment_column: 情感标签列名
        click_column: 点击数列名
        comment_column: 评论数列名
        majia_column: 马甲账号标识列名
        click_weight: 点击量权重系数
        comment_weight: 评论数权重系数
        majia_penalty: 马甲账号的权重惩罚系数
        normalize: 是否将指数归一化到[-100, 100]区间
    
    返回:
        按日期索引的情感指数DataFrame
    """
    # 确保日期列为datetime类型
    if date_column in df.columns:
        df[date_column] = pd.to_datetime(df[date_column])
    else:
        raise ValueError(f"日期列'{date_column}'不存在")
    
    # 确保情感列存在
    if sentiment_column not in df.columns:
        raise ValueError(f"情感列'{sentiment_column}'不存在")
    
    # 准备计算权重所需的列
    required_columns = [date_column, sentiment_column]
    for col in [click_column, comment_column, majia_column]:
        if col not in df.columns:
            print(f"警告：列'{col}'不存在，将不会用于权重计算")
    
    # 创建工作副本，避免修改原始数据
    working_df = df.copy()
    
    # 将情感标签映射为数值
    working_df['sentiment_value'] = working_df[sentiment_column].apply(map_sentiment_to_value)
    
    # 计算每条评论的权重
    working_df['weight'] = working_df.apply(
        lambda row: calculate_weight(
            row, 
            click_weight=click_weight, 
            comment_weight=comment_weight, 
            majia_penalty=majia_penalty
        ), 
        axis=1
    )
    
    # 计算加权情感得分
    working_df['weighted_sentiment'] = working_df['sentiment_value'] * working_df['weight']
    
    # 按日期分组计算
    working_df['date'] = working_df[date_column].dt.date
    
    # 按日期分组并计算加权平均情感指数
    daily_sentiment = working_df.groupby('date').agg(
        total_weighted_sentiment=('weighted_sentiment', 'sum'),
        total_weight=('weight', 'sum'),
        comment_count=('sentiment_value', 'count'),
        positive_count=('sentiment_value', lambda x: sum(x > 0)),
        neutral_count=('sentiment_value', lambda x: sum(x == 0)),
        negative_count=('sentiment_value', lambda x: sum(x < 0))
    )
    
    # 计算最终的情感指数
    daily_sentiment['sentiment_index'] = daily_sentiment['total_weighted_sentiment'] / daily_sentiment['total_weight']
    
    # 计算情感比例
    daily_sentiment['positive_ratio'] = daily_sentiment['positive_count'] / daily_sentiment['comment_count']
    daily_sentiment['neutral_ratio'] = daily_sentiment['neutral_count'] / daily_sentiment['comment_count']
    daily_sentiment['negative_ratio'] = daily_sentiment['negative_count'] / daily_sentiment['comment_count']
    
    # 归一化到[-100, 100]区间
    if normalize:
        daily_sentiment['sentiment_index'] = daily_sentiment['sentiment_index'] * 100
    
    return daily_sentiment


def plot_sentiment_index(daily_sentiment, output_file=None, title="股票情感指数走势", figsize=(12, 8)):
    """
    绘制情感指数走势图
    
    参数:
        daily_sentiment: 按日期索引的情感指数DataFrame
        output_file: 输出文件路径，如果为None则仅显示不保存
        title: 图表标题
        figsize: 图表大小
    
    返回:
        None
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize, sharex=True, gridspec_kw={'height_ratios': [3, 1]})
    
    # 绘制情感指数曲线
    dates = daily_sentiment.index
    index_values = daily_sentiment['sentiment_index']
    
    ax1.plot(dates, index_values, 'b-', linewidth=2)
    ax1.set_title(title, fontsize=16)
    ax1.set_ylabel('情感指数 [-100, 100]', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.axhline(y=0, color='r', linestyle='-', alpha=0.3)
    
    # 填充正负区域
    ax1.fill_between(dates, index_values, 0, where=(index_values >= 0), color='green', alpha=0.3, interpolate=True)
    ax1.fill_between(dates, index_values, 0, where=(index_values < 0), color='red', alpha=0.3, interpolate=True)
    
    # 添加情感比例堆叠条形图
    positive_ratio = daily_sentiment['positive_ratio']
    neutral_ratio = daily_sentiment['neutral_ratio']
    negative_ratio = daily_sentiment['negative_ratio']
    
    ax2.bar(dates, positive_ratio, color='green', alpha=0.7, label='正面')
    ax2.bar(dates, neutral_ratio, bottom=positive_ratio, color='gray', alpha=0.7, label='中性')
    ax2.bar(dates, negative_ratio, bottom=positive_ratio+neutral_ratio, color='red', alpha=0.7, label='负面')
    
    ax2.set_ylabel('情感比例', fontsize=12)
    ax2.set_xlabel('日期', fontsize=12)
    ax2.grid(True, linestyle='--', alpha=0.3)
    ax2.legend(loc='upper right')
    
    plt.tight_layout()
    
    # 如果指定了输出文件，则保存图表
    if output_file:
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"情感指数走势图已保存到: {output_file}")
    
    # 显示图表
    plt.show()


def export_sentiment_index(daily_sentiment, output_file):
    """
    导出情感指数数据到CSV文件
    
    参数:
        daily_sentiment: 按日期索引的情感指数DataFrame
        output_file: 输出文件路径
    
    返回:
        None
    """
    # 确保输出目录存在
    output_dir = os.path.dirname(output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 导出到CSV
    daily_sentiment.to_csv(output_file)
    print(f"情感指数数据已导出到: {output_file}")


def main():
    """示例用法"""
    import argparse
    
    parser = argparse.ArgumentParser(description="计算股票评论情感指数")
    parser.add_argument('--input', type=str, required=True, help='输入CSV文件路径')
    parser.add_argument('--output', type=str, default=None, help='输出CSV文件路径')
    parser.add_argument('--date_column', type=str, default='publish_date', help='日期列名称')
    parser.add_argument('--sentiment_column', type=str, default='sentiment', help='情感列名称')
    parser.add_argument('--click_column', type=str, default='post_click_count', help='点击数列名称')
    parser.add_argument('--comment_column', type=str, default='post_comment_count', help='评论数列名称')
    parser.add_argument('--majia_column', type=str, default='user_is_majia', help='马甲标识列名称')
    parser.add_argument('--click_weight', type=float, default=1.0, help='点击量权重系数')
    parser.add_argument('--comment_weight', type=float, default=2.0, help='评论数权重系数')
    parser.add_argument('--majia_penalty', type=float, default=0.5, help='马甲账号权重惩罚系数')
    parser.add_argument('--plot', action='store_true', help='是否生成图表')
    parser.add_argument('--plot_output', type=str, default=None, help='图表输出文件路径')
    parser.add_argument('--title', type=str, default="股票情感指数走势", help='图表标题')
    
    args = parser.parse_args()
    
    # 读取CSV文件
    df = pd.read_csv(args.input)
    
    # 计算情感指数
    daily_sentiment = calculate_daily_sentiment_index(
        df,
        date_column=args.date_column,
        sentiment_column=args.sentiment_column,
        click_column=args.click_column,
        comment_column=args.comment_column,
        majia_column=args.majia_column,
        click_weight=args.click_weight,
        comment_weight=args.comment_weight,
        majia_penalty=args.majia_penalty
    )
    
    # 输出结果
    print("\n每日情感指数:")
    print(daily_sentiment)
    
    # 导出到CSV
    if args.output:
        export_sentiment_index(daily_sentiment, args.output)
    
    # 绘制图表
    if args.plot:
         plot_sentiment_index(daily_sentiment, output_file=args.plot_output, title=args.title)
    
    return daily_sentiment


if __name__ == "__main__":
    main() 